#!/usr/bin/env python3
# Copyright 2019-2022 Alibaba Group Holding Limited.
# SPDX-License-Identifier: GPL-2.0 OR BSD-3-Clause

import subprocess
import sh
import yaml
import os
import sys
from glob import glob

class TestPublicVar:
    def setup_class(self):
        print("Public vars test")
        cur_sys = str(sh.uname('-r')).strip()
        scheduler_rpm = glob(os.path.join('/tmp/work', 'scheduler*.rpm'))
        if len(scheduler_rpm) != 1:
            print("Please check your scheduler rpm");
            self.teardown_class()
            sys.exit(1)
        scheduler_rpm = scheduler_rpm[0]
        sh.rpm('-ivh', scheduler_rpm)
        module = '/var/plugsched/' + cur_sys + '/scheduler.ko'
        yaml_file = '/tmp/work/scheduler/working/boundary.yaml'
        with open(yaml_file, 'r') as f:
            yaml_version = float(yaml.__version__)
            if yaml_version >= 5.1:
                config = yaml.load(f, Loader=yaml.FullLoader)
            else:
                config = yaml.load(f)
            self.public_vars = config['global_var']['extra_public']
        cmd = "objdump -t " + module
        self.symtab = str(subprocess.check_output(cmd, shell=True)).split('\\n')

    def test_syms(self):
        for var in self.public_vars:
            for cur_line in self.symtab:
                if cur_line == '' or cur_line.split()[-1] != var:
                    continue
                if not '*UND*' in cur_line:
                    self.error_handler(var)
                break

    def error_handler(self, var):
        print("Public var: " + str(var) + "is not UND")
        self.teardown_class()
        sys.exit(1)

    def teardown_class(self):
        tmp = subprocess.Popen("lsmod | grep scheduler", shell=True, stdout=subprocess.PIPE)
        if tmp.stdout.read() != b'':
            sh.rpm('-e', 'scheduler-xxx')

if __name__ == '__main__':
    test_unit = TestPublicVar()
    test_unit.setup_class()
    test_unit.test_syms()
    test_unit.teardown_class()
